        self.tools = StockTools()

"""
基于 AutoGen 的股票分析智能体
"""
import autogen
from autogen import AssistantAgent, UserProxyAgent, ConversableAgent
import json
import os
from typing import Dict, Any, List, Optional
from colorama import Fore, Style
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from tools.stock_tools import StockTools


from agents.business_agents import BusinessAgentFactory


class StockAgentFactory(BusinessAgentFactory):
    """股票分析智能体工厂类"""
    
    def __init__(self, config_path: str = None):
        """
        初始化智能体工厂
        
        Args:
            config_path: 配置文件路径
        """
        self.config = self._load_config(config_path)
        self.llm_config = self.config["llm_config"]
        self.agents_config = self.config["agents"]
        BusinessAgentFactory.__init__(self, self.llm_config)
        self.tools = StockTools()
        
    def _load_config(self, config_path: str = None) -> Dict[str, Any]:
        """加载配置文件"""
        if config_path is None:
            config_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), "config.json")
        
        try:
            with open(config_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(Fore.RED + f"加载配置文件失败: {e}")
            raise
    
    def create_data_collector(self) -> AssistantAgent:
        """创建数据收集智能体"""
        agent_config = self.agents_config["data_collector"]
        
        # 定义函数映射，用于工具调用
        function_map = {
            "get_stock_data": self.tools.get_stock_data,
            "search_stock_news": self.tools.search_stock_news,
        }
        
        # 创建智能体
        agent = AssistantAgent(
            name=agent_config["name"],
            system_message=agent_config["system_message"] + """
            
你可以使用以下工具：
1. get_stock_data(stock_code, days=365) - 获取股票历史数据
2. search_stock_news(stock_code, query=None) - 搜索股票相关新闻

在执行任务时，请：
1. 首先使用 get_stock_data 获取股票数据
2. 然后使用 search_stock_news 获取相关新闻
3. 整理收集到的信息，生成结构化报告
""",
            llm_config=self.llm_config,
            max_consecutive_auto_reply=3,
        )
        
        # 注册函数
        for func_name, func in function_map.items():
            agent.register_for_execution(name=func_name)(func)
        
        print(Fore.GREEN + f"✅ 创建智能体: {agent_config['name']}")
        return agent
    
    def create_data_analyst(self, stock_code: str) -> AssistantAgent:
        """创建数据分析智能体"""
        agent_config = self.agents_config["data_analyst"]
        
        # 定义函数映射
        function_map = {
            "analyze_stock_data": lambda stock_data, analysis_type='comprehensive': self.tools.analyze_stock_data(stock_code, stock_data, analysis_type),
        }
        
        # 创建智能体
        agent = AssistantAgent(
            name=agent_config["name"],
            system_message=agent_config["system_message"] + """
            
你可以使用以下工具：
1. analyze_stock_data(stock_data, analysis_type='comprehensive') - 分析股票数据

在分析时，请：
1. 进行技术分析：分析价格趋势、成交量、技术指标
2. 进行基本面分析：评估公司财务、行业地位、估值水平
3. 进行风险评估：识别潜在风险、评估风险等级
4. 生成综合分析报告，包含关键发现和投资建议
""",
            llm_config=self.llm_config,
            max_consecutive_auto_reply=3,
        )
        
        # 注册函数
        for func_name, func in function_map.items():
            agent.register_for_execution(name=func_name)(func)
        
        print(Fore.GREEN + f"✅ 创建智能体: {agent_config['name']}")
        return agent
    
    def create_decision_maker(self) -> AssistantAgent:
        """创建投资决策智能体"""
        agent_config = self.agents_config["decision_maker"]
        
        # 创建智能体
        agent = AssistantAgent(
            name=agent_config["name"],
            system_message=agent_config["system_message"] + """
            
基于收到的数据和分析报告，请：
1. 综合评估投资价值
2. 制定投资策略（买入/持有/卖出）
3. 设定目标价格和止损位
4. 提供具体的操作建议和时间安排
5. 强调主要风险和注意事项

你的决策应该：
- 明确具体，可执行
- 包含量化指标
- 考虑风险收益比
- 提供清晰的理由
""",
            llm_config=self.llm_config,
            max_consecutive_auto_reply=3,
        )
        
        print(Fore.GREEN + f"✅ 创建智能体: {agent_config['name']}")
        return agent
    
    def create_user_proxy(self) -> UserProxyAgent:
        """创建用户代理智能体"""
        agent_config = self.agents_config["user_proxy"]
        
        # 创建用户代理
        agent = UserProxyAgent(
            name=agent_config["name"],
            system_message=agent_config["system_message"],
            human_input_mode="NEVER",  # 不需要人工输入
            max_consecutive_auto_reply=0,  # 不自动回复
            code_execution_config=False,  # 不执行代码
        )
        
        print(Fore.GREEN + f"✅ 创建智能体: {agent_config['name']}")
        return agent
    
    def create_all_agents(self, stock_code: str) -> Dict[str, ConversableAgent]:
        """创建所有智能体"""
        agents = {
            "data_collector": self.create_data_collector(),
            "data_analyst": self.create_data_analyst(stock_code),
            "decision_maker": self.create_decision_maker(),
            "user_proxy": self.create_user_proxy(),
            "market_researcher": self.create_market_research_agent(),
            "risk_manager": self.create_risk_management_agent(),
            "portfolio_manager": self.create_portfolio_management_agent(),
        }
        
        print(Fore.CYAN + f"\n📊 成功创建 {len(agents)} 个智能体")
        return agents


class StockAnalysisTeam:
    """股票分析团队（智能体协作）"""
    
    def __init__(self, agents: Dict[str, ConversableAgent]):
        """
        初始化分析团队
        
        Args:
            agents: 智能体字典
        """
        self.agents = agents
        self.data_collector = agents["data_collector"]
        self.data_analyst = agents["data_analyst"]
        self.decision_maker = agents["decision_maker"]
        self.user_proxy = agents["user_proxy"]
        
    def analyze_stock_sequential(self, stock_code: str) -> str:
        """
        顺序执行股票分析
        
        Args:
            stock_code: 股票代码
            
        Returns:
            分析结果
        """
        print(Fore.CYAN + Style.BRIGHT + f"\n🚀 开始分析股票: {stock_code}")
        
        results = {}
        
        # Step 1: 数据收集
        print(Fore.YELLOW + "\n📊 步骤 1/3: 数据收集")
        collection_task = f"""
请收集股票 {stock_code} 的相关数据：
1. 获取最近一年的历史价格数据
2. 搜索最新的相关新闻
3. 整理数据生成报告
"""
        self.user_proxy.initiate_chat(
            self.data_collector,
            message=collection_task,
            max_turns=2,
            summary_method="reflection_with_llm"
        )
        results["collection"] = self.user_proxy.last_message()["content"]
        
        # Step 2: 数据分析
        print(Fore.YELLOW + "\n📈 步骤 2/3: 数据分析")
        analysis_task = f"""
基于以下收集到的数据，请进行深入分析：

{results['collection'][:2000]}

请进行：
1. 技术分析
2. 基本面分析
3. 风险评估
4. 生成分析报告
"""
        self.user_proxy.initiate_chat(
            self.data_analyst,
            message=analysis_task,
            max_turns=2,
            summary_method="reflection_with_llm"
        )
        results["analysis"] = self.user_proxy.last_message()["content"]
        
        # Step 3: 投资决策
        print(Fore.YELLOW + "\n🎯 步骤 3/3: 投资决策")
        decision_task = f"""
基于以下分析报告，请制定投资决策：

分析摘要：
{results['analysis'][:2000]}

请提供：
1. 投资建议（强烈买入/买入/持有/卖出/强烈卖出）
2. 目标价格
3. 止损位
4. 操作策略
5. 风险提示
"""
        self.user_proxy.initiate_chat(
            self.decision_maker,
            message=decision_task,
            max_turns=2,
            summary_method="reflection_with_llm"
        )
        results["decision"] = self.user_proxy.last_message()["content"]
        
        # 生成最终报告
        final_report = self._generate_final_report(stock_code, results)
        return final_report
    
    def _generate_final_report(self, stock_code: str, results: Dict[str, str]) -> str:
        """生成最终报告"""
        from datetime import datetime
        
        report = f"""
{'=' * 70}
股票分析报告 - {stock_code}
{'=' * 70}

生成时间: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}
分析框架: AutoGen Multi-Agent System

一、数据收集结果
{'-' * 50}
{results.get('collection', 'N/A')[:1500]}

二、深度分析结果
{'-' * 50}
{results.get('analysis', 'N/A')[:1500]}

三、投资决策建议
{'-' * 50}
{results.get('decision', 'N/A')[:1500]}

{'=' * 70}
免责声明：本报告仅供参考，不构成投资建议。投资有风险，决策需谨慎。
{'=' * 70}
"""
        return report
